This is the code repository of the AISTATS 2022 paper : "Predicting the impact of treatments over time with uncertainty awareneural differential equations."
We support both conda environments and poetry types of environments.
conda create --name causalode --file requirements.txt
Poetry is a python dependencies manager. You can install it following the instructions here.
You can then install all the dependencies at once by running:
poetry install
Attention When using poetry as package dependencies, you'll have to add poetry run
to your commands. Example:
poetry run python train_model.py
To run the experiments from the paper, you can use the parametrizations below. All datasets are generated on the fly. The different folds were created by changing the seed parameter.
Harmonic Oscillator data set
python train_model.py --N_ts=1000 --batch_size=64 --continuous_treatment=True --embedding_dim=64 --gamma=8 --kl_param=0.001 --logger_type=wandb --max_epochs=2000 --num_samples=5 --propensity_scores=False --seed=44 --sigma_sde=0.1 --static=False --std_dev=False --variational=True
CardioVascular data set
python train_model.py --N_ts=1000 --batch_size=64 --continuous_treatment=True --dataset_name=cv --embedding_dim=64 --gamma=10 --horizon=10 --kl_param=0.001 --logger_type=wandb --max_epochs=2000 --num_samples=5 --propensity_scores=False --seed=44 --sigma_sde=0.1 --std_dev=False --variational=True
Dexamethasone data set
train_model.py --N_ts=1000 --batch_size=64 --continuous_treatment=True --dataset_name=covid --embedding_dim=64 --gamma=10 --horizon=10 --kl_param=0.001 --logger_type=wandb --max_epochs=2000 --num_samples=5 --propensity_scores=False --seed=43 --sigma_sde=0.1 --std_dev=False --variational=True
We provide jupyter notebooks to process the results of the experiments and generate the figures displayed in the paper.
computes the results reported in Table 1 in the paper.
creates the Figure 3 in the paper.
contains the experiments of section 5.3 about improving the treatment strategies with the derived uncertainties
contains the scripts for the experiment on out-of-distribution sampels as reported in the Appendix.
All experiments have been logged with Weights and Biases but we also support MLFlow and Tensorboard. You change the type of logger by setting the --logger_type
option.
This project welcomes contributions and suggestions. Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.
When you submit a pull request, a CLA bot will automatically determine whether you need to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions provided by the bot. You will only need to do this once across all repos using our CLA.
This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact opencode@microsoft.com with any additional questions or comments.
This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft trademarks or logos is subject to and must follow Microsoft's Trademark & Brand Guidelines. Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship. Any use of third-party trademarks or logos are subject to those third-party's policies.